import os
import struct
from pathlib import Path

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding, hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

from interface.crypto_interfaces import CryptoProcessor


class AESCryptoProcessor(CryptoProcessor):
    """通用文件格式的AES-256-CBC加密处理器"""

    def __init__(self, password: str):
        self.password = password
        self.salt_length = 16
        self.iv_length = 16
        # 定义特殊标记，用于标识加密数据的开始
        self.MAGIC_MARKER = b"SECRET_START"
        self.MARKER_LENGTH = len(self.MAGIC_MARKER)
        self.HEADER_LENGTH = self.MARKER_LENGTH + self.salt_length + self.iv_length + 8

    def _derive_key(self, salt: bytes) -> bytes:
        """派生加密密钥"""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        return kdf.derive(self.password.encode())

    def merge_and_encrypt(self, cover_file: Path, secret_file: Path, output_path: Path):
        """
        将密文文件加密后嵌入到载体文件中

        :param cover_file: 载体文件路径
        :param secret_file: 需要加密隐藏的文件路径
        :param output_path: 输出文件路径（可以是目录或文件）
        """
        try:
            # 读取原始文件和待加密文件
            with open(cover_file, "rb") as f:
                cover_data = f.read()
            with open(secret_file, "rb") as f:
                secret_data = f.read()

            self.log.info("文件读取完成")

            # 生成加密参数
            salt = os.urandom(self.salt_length)
            iv = os.urandom(self.iv_length)
            key = self._derive_key(salt)

            # 加密数据
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
            encryptor = cipher.encryptor()
            padder = padding.PKCS7(128).padder()
            padded_data = padder.update(secret_data) + padder.finalize()
            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
            self.log.info("加密完成")

            # 确保输出路径是文件
            if output_path.is_dir():
                output_path = output_path / cover_file.name

            # 计算加密数据长度
            encrypted_length = len(encrypted_data)

            # 合并文件：载体内容 + 特殊标记 + 加密数据长度 + salt + iv + 加密数据
            with open(output_path, "wb") as f:
                # 写入载体内容
                f.write(cover_data)

                # 写入特殊标记
                f.write(self.MAGIC_MARKER)

                # 写入加密数据长度（8字节大端序）
                f.write(struct.pack('>Q', encrypted_length))

                # 写入加密参数
                f.write(salt)
                f.write(iv)

                # 写入加密数据
                f.write(encrypted_data)

            self.log.info(f"合并文件保存完成: {output_path}")
            return output_path
        except FileNotFoundError as e:
            self.log.error(f"文件未找到: {e}")
            raise
        except ValueError as e:
            self.log.error(f"加密错误: {e}")
            raise
        except Exception as e:
            self.log.critical(f"未知错误: {e}", exc_info=True)
            raise

    def decrypt_and_extract(self, embedded_file: Path, output_path: Path):
        """
        从嵌入文件中提取和解密密文文件

        :param embedded_file: 包含加密数据的文件
        :param output_path: 解密后文件的输出路径
        """
        try:
            # 使用流式处理，避免大文件内存问题
            with open(embedded_file, "rb") as f:
                # 查找特殊标记
                marker_found = False
                buffer = b""
                chunk_size = 4096

                # 扫描整个文件查找特殊标记
                while True:
                    chunk = f.read(chunk_size)
                    if not chunk:
                        break

                    buffer += chunk
                    marker_pos = buffer.find(self.MAGIC_MARKER)

                    # 如果找到标记
                    if marker_pos != -1:
                        # 计算标记在文件中的位置
                        file_pos = f.tell() - len(buffer) + marker_pos
                        marker_found = True
                        break

                    # 保留部分数据以防标记跨块
                    keep = min(len(buffer), len(self.MAGIC_MARKER) - 1)
                    buffer = buffer[-keep:]

            if not marker_found:
                raise ValueError("未找到加密数据标记")

            self.log.info(f"找到加密数据标记于位置: {file_pos}")

            # 定位到标记位置
            with open(embedded_file, "rb") as f:
                f.seek(file_pos + self.MARKER_LENGTH)

                # 读取加密数据长度
                encrypted_length_bytes = f.read(8)
                if len(encrypted_length_bytes) != 8:
                    raise ValueError("无法读取加密数据长度")
                encrypted_length = struct.unpack('>Q', encrypted_length_bytes)[0]

                # 读取加密参数
                salt = f.read(self.salt_length)
                if len(salt) != self.salt_length:
                    raise ValueError("无法读取salt值")

                iv = f.read(self.iv_length)
                if len(iv) != self.iv_length:
                    raise ValueError("无法读取iv值")

                # 读取加密数据
                encrypted_data = f.read(encrypted_length)
                if len(encrypted_data) != encrypted_length:
                    raise ValueError(f"加密数据长度不匹配: 预期 {encrypted_length}, 实际 {len(encrypted_data)}")

            self.log.info(f"读取加密数据完成: {len(encrypted_data)}字节")

            # 派生密钥
            key = self._derive_key(salt)

            # 解密数据
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
            decryptor = cipher.decryptor()
            unpadder = padding.PKCS7(128).unpadder()

            decrypted_padded = decryptor.update(encrypted_data) + decryptor.finalize()
            decrypted_data = unpadder.update(decrypted_padded) + unpadder.finalize()
            self.log.info(f"解密完成: {len(decrypted_data)}字节")

            # 确保输出目录存在
            output_path.parent.mkdir(parents=True, exist_ok=True)

            # 保存解密文件
            with open(output_path, "wb") as f:
                f.write(decrypted_data)

            self.log.info(f"解密文件保存完成: {output_path}")
            return output_path
        except FileNotFoundError as e:
            self.log.error(f"文件未找到: {e}")
            raise
        except ValueError as e:
            self.log.error(f"解密错误: {e}")
            raise
        except Exception as e:
            self.log.critical(f"未知错误: {e}", exc_info=True)
            raise


if __name__ == '__main__':
    processor = AESCryptoProcessor(password="123456")
    image_paths = Path(r"E:\下载文件\联想下载\1b87e2c5880511ebb6edd017c2d2eca2.png")
    archive_paths = Path(r"E:\下载文件\新建文件夹\魔法.7删除z")
    output_paths = Path(r"E:\下载文件\新建文件夹")
    merged_paths = output_paths / image_paths.name
    processor.merge_and_encrypt(
        cover_file=image_paths,
        secret_file=archive_paths,
        output_path=output_paths
    )

    processor.decrypt_and_extract(embedded_file=merged_paths, output_path=Path(r"E:\下载文件\新建文件夹\解密.zip"))
